#include <vector>
using namespace std;
#include <boost/shared_ptr.hpp>
#include <wx/wx.h>
#include <wx/string.h>
#include <wx/dir.h>
#include <wx/filename.h>
#include <wx/listctrl.h>
#include "ItemInfo.h"
#include "Utils.h"
#include <shlobj.h>


wxDirTraverserMy::wxDirTraverserMy(wxArrayString& pFiles, wxArrayString& pDirectories, vector<LARGE_INTEGER> &pFileSizes, LARGE_INTEGER &pTotalFileSize, bool pCalcualteSizes) : files(pFiles), directories(pDirectories), fileSizes(pFileSizes), totalFileSize(pTotalFileSize), calcualteSizes(pCalcualteSizes)
{
    totalFileSize.QuadPart = 0;
}

wxDirTraverseResult wxDirTraverserMy::OnFile(const wxString& filename)
{
    DWORD attributes = GetFileAttributes(filename);
    if (attributes & FILE_ATTRIBUTE_HIDDEN || attributes & FILE_ATTRIBUTE_SYSTEM)
        return wxDIR_CONTINUE;
    files.Add(filename);
    if (calcualteSizes)
    {
        LARGE_INTEGER fileSize;
        HANDLE h = CreateFile(filename, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, 0);
        if (!GetFileSizeEx(h, &fileSize))
            fileSize.QuadPart = 0;
        CloseHandle(h);
        fileSizes.push_back(fileSize);
        totalFileSize.QuadPart += fileSize.QuadPart;
    }
    return wxDIR_CONTINUE;
}

wxDirTraverseResult wxDirTraverserMy::OnDir(const wxString& dirname)
{
    DWORD attributes = GetFileAttributes(dirname);
    if (attributes & FILE_ATTRIBUTE_HIDDEN || attributes & FILE_ATTRIBUTE_SYSTEM)
        return wxDIR_IGNORE;
    directories.Add(dirname);
    return wxDIR_IGNORE;
}


bool GetDirectoryList(const wxString &pPath, vector<ItemInfo *> &list)
{
    if (!wxFileName::IsDirReadable(pPath))
        return false;
    wxString path = pPath;
    if (path.Last() == '\\')
        path.RemoveLast();
    wxArrayString files;
    wxArrayString directories;
    LARGE_INTEGER fileSize;
    vector<LARGE_INTEGER> fileSizes;
    wxDirTraverserMy traverser(files, directories, fileSizes, fileSize, true);
    wxDir dir(path);

    dir.Traverse(traverser);

    wxString topPath, topPathName;
    wxFileName::SplitPath(path, &topPath, &topPathName, NULL, wxPATH_NATIVE);
    if (!topPathName.IsEmpty())
    {
        ItemInfo *topInfo = new ItemInfo;

        topInfo->isDirectory = true;
        topInfo->fullPath = topPath;
        topInfo->name = "..";
        topInfo->icon = GetIconFromFile(topInfo->fullPath, topInfo->isDirectory);
        list.push_back(topInfo);
    }

    for (unsigned int i = 0; i < directories.size(); i++)
    {
        wxString dirName;
        wxFileName::SplitPath(directories[i], NULL, &dirName, NULL, wxPATH_NATIVE);

        ItemInfo *info = new ItemInfo;
        info->isDirectory = true;
        info->name = dirName;
        info->fullPath = directories[i];
        info->icon = GetIconFromFile(info->fullPath, info->isDirectory);

        list.push_back(info);
    }
    for (unsigned int i = 0; i < files.size(); i++)
    {
        wxFileName filename(files[i]);

        ItemInfo *info = new ItemInfo;
        info->isDirectory = false;
        info->fullPath = files[i];
        info->name = filename.GetFullName();
        //info->icon = GetIconFromFile(info->fullPath, info->isDirectory);
        info->size = fileSizes[i].QuadPart;

        list.push_back(info);
    }
    return true;
}

bool PopulateListFromDirectory(wxListView *listView, const wxString &pPath)
{
    if (!wxFileName::IsDirReadable(pPath))
        return false;
    wxString path = pPath;
    wxArrayString files;
    wxArrayString directories;
    LARGE_INTEGER fileSize;
    vector<LARGE_INTEGER> fileSizes;
    wxDirTraverserMy traverser(files, directories, fileSizes, fileSize, false);
    wxDir dir(path);
    listView->Freeze();
    for (int i = 0; i < listView->GetItemCount(); i++)
    {
        delete (void*)listView->GetItemData(i);
    }

    listView->DeleteAllItems();
    dir.Traverse(traverser);
    long id = 0;

    wxString topPath, topPathName;
    wxFileName::SplitPath(path, &topPath, &topPathName, NULL, wxPATH_NATIVE);
    if (!topPathName.IsEmpty())
    {
        ItemInfo *topInfo = new ItemInfo;

        topInfo->isDirectory = true;
        topInfo->fullPath = topPath;
        topInfo->name = "..";
        wxListItem topItem;
        topItem.SetText(topInfo->name);
        topItem.SetData(topInfo);
        topItem.SetColumn(0);
        topItem.SetId(id);
        listView->InsertItem(topItem);
        id++;
    }

    for (unsigned int i = 0; i < directories.size(); i++)
    {
        wxString dirName;
        wxFileName::SplitPath(directories[i], NULL, &dirName, NULL, wxPATH_NATIVE);

        ItemInfo *info = new ItemInfo;
        info->isDirectory = true;
        info->name = dirName;
        info->fullPath = directories[i];

        wxListItem item;
        item.SetText(info->name);
        item.SetData(info);
        item.SetColumn(0);
        item.SetId(id);
        wxIcon icon = GetIconFromFile(info->fullPath, true);
        int iconIndex = listView->GetImageList(wxIMAGE_LIST_SMALL)->Add(icon);
        item.SetImage(iconIndex);
        listView->InsertItem(item);
        id++;
    }
    for (unsigned int i = 0; i < files.size(); i++)
    {
        wxFileName filename(files[i]);

        ItemInfo *info = new ItemInfo;
        info->isDirectory = false;
        info->fullPath = files[i];
        info->name = filename.GetFullName();

        wxListItem item;
        item.SetText(info->name);
        item.SetData(info);
        item.SetColumn(0);
        item.SetId(id);
        wxIcon icon = GetIconFromFile(info->fullPath);
        int iconIndex = listView->GetImageList(wxIMAGE_LIST_SMALL)->Add(icon);
        item.SetImage(iconIndex);
        listView->InsertItem(item);
        id++;
    }
    listView->Thaw();
    return true;
}

HRESULT GetUIObjectOfFile(HWND hwnd, LPCWSTR pszPath, REFIID riid, void **ppv)
{
    *ppv = NULL;
    HRESULT hr;
    LPITEMIDLIST pidl;
    SFGAOF sfgao;
    if (SUCCEEDED(hr = SHParseDisplayName(pszPath, NULL, &pidl, 0, &sfgao))) {
        IShellFolder *psf;
        LPCITEMIDLIST pidlChild;
        if (SUCCEEDED(hr = SHBindToParent(pidl, IID_IShellFolder,
                                          (void**)&psf, &pidlChild))) {
            hr = psf->GetUIObjectOf(hwnd, 1, &pidlChild, riid, NULL, ppv);
            psf->Release();
        }
        CoTaskMemFree(pidl);
    }
    return hr;
}

#define SCRATCH_QCM_FIRST 1
#define SCRATCH_QCM_LAST  0x7FFF
IContextMenu2 *g_pcm2 = NULL;
IContextMenu3 *g_pcm3 = NULL;

void DisplayContextMenu(HWND hwnd, HWND hwndContext, LPCWSTR pszPath, UINT xPos, UINT yPos)
{
    POINT pt = { xPos, yPos };

    IContextMenu *pcm;
    if (SUCCEEDED(GetUIObjectOfFile(hwnd, pszPath, IID_IContextMenu, (void**)&pcm)))
    {
        HMENU hmenu = CreatePopupMenu();
        if (hmenu)
        {
            if (SUCCEEDED(pcm->QueryContextMenu(hmenu, 0, SCRATCH_QCM_FIRST, SCRATCH_QCM_LAST, CMF_NORMAL)))
            {
                pcm->QueryInterface(IID_IContextMenu2, (void**)&g_pcm2);
                pcm->QueryInterface(IID_IContextMenu3, (void**)&g_pcm3);
                int iCmd = TrackPopupMenuEx(hmenu, TPM_RETURNCMD, pt.x, pt.y, hwnd, NULL);
                if (g_pcm2)
                {
                    g_pcm2->Release();
                    g_pcm2 = NULL;
                }
                if (g_pcm3)
                {
                    g_pcm3->Release();
                    g_pcm3 = NULL;
                }
                if (iCmd > 0)
                {
                    CMINVOKECOMMANDINFOEX info = { 0 };
                    info.cbSize = sizeof(info);
                    info.fMask = CMIC_MASK_UNICODE | CMIC_MASK_PTINVOKE;
                    if (GetKeyState(VK_CONTROL) < 0)
                    {
                        info.fMask |= CMIC_MASK_CONTROL_DOWN;
                    }
                    if (GetKeyState(VK_SHIFT) < 0)
                    {
                        info.fMask |= CMIC_MASK_SHIFT_DOWN;
                    }
                    info.hwnd = hwnd;
                    info.lpVerb  = MAKEINTRESOURCEA(iCmd - SCRATCH_QCM_FIRST);
                    info.lpVerbW = MAKEINTRESOURCEW(iCmd - SCRATCH_QCM_FIRST);
                    info.nShow = SW_SHOWNORMAL;
                    info.ptInvoke = pt;
                    pcm->InvokeCommand((LPCMINVOKECOMMANDINFO)&info);
                }
            }
            DestroyMenu(hmenu);
        }
        pcm->Release();
    }
}

void InvokeDefaultItem(HWND hwnd, HWND hwndContext, LPCWSTR pszPath, UINT xPos, UINT yPos)
{
    IContextMenu *pcm;
    if (SUCCEEDED(GetUIObjectOfFile(hwnd, pszPath, IID_IContextMenu, (void**)&pcm)))
    {
        HMENU hmenu = CreatePopupMenu();
        if (hmenu)
        {
            if (SUCCEEDED(pcm->QueryContextMenu(hmenu, 0, SCRATCH_QCM_FIRST, SCRATCH_QCM_LAST, CMF_DEFAULTONLY)))
            {
                UINT id = GetMenuDefaultItem(hmenu, FALSE, 0);
                if (id != (UINT) - 1)
                {
                    CMINVOKECOMMANDINFO info = { 0 };
                    info.cbSize = sizeof(info);
                    info.hwnd = hwnd;
                    info.lpVerb = MAKEINTRESOURCEA(id - SCRATCH_QCM_FIRST);
                    info.nShow = SW_SHOWNORMAL;
                    pcm->InvokeCommand(&info);
                }
            }
            DestroyMenu(hmenu);
        }
        pcm->Release();
    }
}

wxIcon GetIconFromFile(const wxString &file, bool isDirectory)
{
    SHFILEINFO stFileInfo;
    DWORD fileAttributes;
    if (isDirectory)
        fileAttributes = FILE_ATTRIBUTE_DIRECTORY;
    else
        fileAttributes = FILE_ATTRIBUTE_NORMAL;

    SHGetFileInfo(file, fileAttributes, &stFileInfo, sizeof( stFileInfo ), SHGFI_ICON | SHGFI_SMALLICON | SHGFI_USEFILEATTRIBUTES);
    wxIcon ret;
    ret.SetHICON(stFileInfo.hIcon);
    ret.SetSize(16, 16);
    return ret;
}


FileTreeBuilder::FileTreeBuilder()
{
    totalFileSize.QuadPart = 0;
    numberOfFiles = 0;
}
LARGE_INTEGER FileTreeBuilder::GetTotalFileSize()
{
    return totalFileSize;
}
void FileTreeBuilder::BuildTree(const wxString &rootPath, bool calculateSizes)
{
    root = FileTreeNodePtr(new FileTreeNode);
    root->fullPath = rootPath;
    root->totalSize.QuadPart = 0;
    BuildTree(root, rootPath, calculateSizes);
}
void FileTreeBuilder::BuildTree(const wxString &rootPath, const wxArrayString &directories, const wxArrayString &files, bool calculateSizes)
{
    root = FileTreeNodePtr(new FileTreeNode);
    root->totalSize.QuadPart = 0;
    root->fullPath = rootPath;
    root->files = files;
    numberOfFiles += files.size();
    if (calculateSizes)
    {
        for (unsigned int i = 0; i < files.size(); i++)
        {
            LARGE_INTEGER fileSize;
            HANDLE h = CreateFile(files[i], GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, 0);
            if (!GetFileSizeEx(h, &fileSize))
                fileSize.QuadPart = 0;
            totalFileSize.QuadPart += fileSize.QuadPart;
            root->totalSize.QuadPart += fileSize.QuadPart;
            root->fileSizes.push_back(fileSize);
            CloseHandle(h);
        }
    }
    for (unsigned int i = 0; i < directories.size(); i++)
    {
        FileTreeBuilder b;
        b.BuildTree(directories[i], calculateSizes);
        totalFileSize.QuadPart += b.GetTotalFileSize().QuadPart;
        root->totalSize.QuadPart += b.GetTotalFileSize().QuadPart;
        numberOfFiles += b.numberOfFiles;
        root->directories.push_back(b.GetRoot());
    }
}
void FileTreeBuilder::BuildTree(FileTreeNodePtr node, const wxString &rootPath, bool calculateSizes)
{
    wxArrayString directories;
    LARGE_INTEGER fileSize;
    wxDirTraverserMy traverser(node->files, directories, node->fileSizes, fileSize, calculateSizes);
    wxDir dir(rootPath);

    dir.Traverse(traverser);
    totalFileSize.QuadPart += fileSize.QuadPart;
    root->totalSize.QuadPart += fileSize.QuadPart;
    numberOfFiles += node->files.size();
    for (unsigned int i = 0; i < directories.size(); i++)
    {
        FileTreeNodePtr child = FileTreeNodePtr(new FileTreeNode);
        child->fullPath = directories[i];
        child->totalSize.QuadPart = 0;
        node->directories.push_back(child);
        BuildTree(child, directories[i], calculateSizes);
        root->totalSize.QuadPart += child->totalSize.QuadPart;
    }
}
FileTreeNodePtr FileTreeBuilder::GetRoot()
{
    return root;
}
unsigned int FileTreeBuilder::GetNumberOfFiles()
{
    return numberOfFiles;
}